import requests
import hashlib
import zipfile
import os
import time

# --- 配置 ---
SYMBOL = "ETHUSDT"
DATA_TYPE = "aggTrades"  # 数据类型, aggTrades, trades等
BASE_URL_PREFIX = f"https://data.binance.vision/data/futures/um/monthly/{DATA_TYPE}/{SYMBOL}/"
START_DATE_STR = "2024-06"
END_DATE_STR = "2025-06"  # 包含此月份

# --- 注意事项 ---
# aggTrades 文件非常大 (zip约600MB, 解压后约3.5GB+)。
# 请确保:
# 1. 目标磁盘有足够的空间。
# 2. 网络连接稳定, 下载可能需要较长时间。脚本中的timeout设置为300秒(5分钟)。

# --- 路径设置 ---
# 请确保以下路径存在, 或者程序将尝试创建它们
# 使用原始字符串 (raw strings) 来处理Windows路径中的反斜杠
BASE_SAVE_PATH = r"F:\personal\binance_agg_trades\\" + SYMBOL  # 建议使用一个更通用的名称, 如 binance_data
ZIP_SAVE_DIR = os.path.join(BASE_SAVE_PATH, "zip")  # 所有zip文件保存到这里
EXTRACT_DIR = BASE_SAVE_PATH  # 解压文件的目录, 例如 F:\personal\binance_data\ETHUSDT\aggTrades\


# --- 辅助函数 (与原脚本相同, 但增加了下载超时时间) ---

def download_file(url, save_path):
    """从URL下载文件到指定路径。针对大文件增加了超时时间。"""
    try:
        print(f"尝试下载: {url}")
        # 确保目标文件夹存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        # 对于大文件, 使用更长的超时时间
        response = requests.get(url, stream=True, timeout=300)
        response.raise_for_status()  # 如果HTTP请求返回了不成功的状态码, 则抛出HTTPError异常

        # 获取文件总大小, 用于显示进度 (可选, 但对大文件友好)
        total_size = int(response.headers.get('content-length', 0))
        downloaded_size = 0

        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
                downloaded_size += len(chunk)
                # 简单的进度条
                done = int(50 * downloaded_size / total_size) if total_size > 0 else 0
                print(
                    f"\r[{'=' * done}{' ' * (50 - done)}] {downloaded_size / 1024 / 1024:.2f} MB / {total_size / 1024 / 1024:.2f} MB",
                    end='')

        print(f"\n成功下载: {save_path}")
        return True
    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 404:
            print(f"\n文件未找到 (404): {url}")
        else:
            print(f"\nHTTP错误, 下载 {url} 失败: {e}")
    except requests.exceptions.RequestException as e:
        print(f"\n下载 {url} 错误: {e}")
    except Exception as e:
        print(f"\n下载 {url} 时发生未知错误: {e}")
    return False


def verify_checksum(zip_file_path, checksum_file_path):
    """对照CHECKSUM文件验证ZIP文件的SHA256校验和。"""
    if not os.path.exists(zip_file_path):
        print(f"用于校验和验证的ZIP文件不存在: {zip_file_path}")
        return False
    if not os.path.exists(checksum_file_path):
        print(f"CHECKSUM文件不存在: {checksum_file_path}")
        return False

    try:
        with open(checksum_file_path, 'r') as f:
            checksum_content = f.read().strip()
        expected_hash = checksum_content.split()[0]

        print(f"正在计算 {os.path.basename(zip_file_path)} 的SHA256校验和, 请稍候...")
        sha256 = hashlib.sha256()
        with open(zip_file_path, 'rb') as f:
            while True:
                data = f.read(65536)  # 以64k块读取
                if not data:
                    break
                sha256.update(data)
        calculated_hash = sha256.hexdigest()

        if calculated_hash == expected_hash:
            print(f"文件 {os.path.basename(zip_file_path)} 校验和验证通过")
            return True
        else:
            print(f"文件 {os.path.basename(zip_file_path)} 校验和不匹配!")
            print(f"  期望值: {expected_hash}")
            print(f"  计算值: {calculated_hash}")
            return False
    except Exception as e:
        print(f"验证 {os.path.basename(zip_file_path)} 校验和时出错: {e}")
        return False


def unzip_file(zip_file_path, extract_to_dir):
    """将ZIP文件解压到指定目录。"""
    try:
        print(f"正在解压 {os.path.basename(zip_file_path)} 到 {extract_to_dir}, 这可能需要一些时间...")
        os.makedirs(extract_to_dir, exist_ok=True)
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to_dir)
        print(f"成功解压: {os.path.basename(zip_file_path)} 到 {extract_to_dir}")
        return True
    except zipfile.BadZipFile:
        print(f"错误: {os.path.basename(zip_file_path)} 是损坏的ZIP文件")
    except Exception as e:
        print(f"解压 {os.path.basename(zip_file_path)} 时出错: {e}")
    return False


# --- 主逻辑 ---
def main():
    # 创建基础保存目录 (如果不存在)
    os.makedirs(ZIP_SAVE_DIR, exist_ok=True)
    os.makedirs(EXTRACT_DIR, exist_ok=True)

    start_year, start_month = map(int, START_DATE_STR.split('-'))
    end_year, end_month = map(int, END_DATE_STR.split('-'))

    current_year = start_year
    current_month = start_month

    while current_year < end_year or (current_year == end_year and current_month <= end_month):
        year_str = str(current_year)
        month_str = f"{current_month:02d}"  # 月份补零

        # 构建文件名和URL
        zip_filename = f"{SYMBOL}-{DATA_TYPE}-{year_str}-{month_str}.zip"
        checksum_filename = f"{zip_filename}.CHECKSUM"

        zip_url = f"{BASE_URL_PREFIX}{zip_filename}"
        checksum_url = f"{BASE_URL_PREFIX}{checksum_filename}"

        # 定义保存路径
        final_zip_save_path = os.path.join(ZIP_SAVE_DIR, zip_filename)
        checksum_save_path = os.path.join(ZIP_SAVE_DIR, checksum_filename)

        print(f"\n--- 开始处理: {SYMBOL} {DATA_TYPE} {year_str}-{month_str} ---")

        # 检查是否已解压
        expected_csv_filename = f"{SYMBOL}-{DATA_TYPE}-{year_str}-{month_str}.csv"
        expected_csv_path = os.path.join(EXTRACT_DIR, expected_csv_filename)

        if os.path.exists(expected_csv_path):
            print(f"数据已存在于 {expected_csv_path}, 跳过 {zip_filename}.")
            continue

        # 如果ZIP文件存在但未解压, 尝试验证并解压
        if os.path.exists(final_zip_save_path) and not os.path.exists(expected_csv_path):
            print(f"ZIP文件 {final_zip_save_path} 已存在但未解压. 尝试验证和解压...")
            if not os.path.exists(checksum_save_path):
                print(f"对应的CHECKSUM文件不存在, 尝试下载...")
                if not download_file(checksum_url, checksum_save_path):
                    print(f"下载CHECKSUM文件失败. 跳过 {zip_filename}.")
                    time.sleep(1)
                    continue

            if verify_checksum(final_zip_save_path, checksum_save_path):
                unzip_file(final_zip_save_path, EXTRACT_DIR)
            else:
                print(f"已存在的ZIP文件 {final_zip_save_path} 校验和失败. 建议删除并重新运行。")

            if os.path.exists(checksum_save_path):
                try:
                    os.remove(checksum_save_path)
                except OSError as e:
                    print(f"无法删除CHECKSUM文件 {checksum_save_path}: {e}")
            continue

        # 1. 下载ZIP文件
        if not download_file(zip_url, final_zip_save_path):
            print(f"下载ZIP文件 {zip_filename} 失败, 跳过.")
            time.sleep(1)
            continue

        # 2. 下载CHECKSUM文件
        if not download_file(checksum_url, checksum_save_path):
            print(f"下载CHECKSUM文件 {checksum_filename} 失败. 无法验证 {zip_filename}.")
            time.sleep(1)
            continue

        # 3. 验证校验和
        if verify_checksum(final_zip_save_path, checksum_save_path):
            # 4. 如果校验通过, 则解压文件
            if not unzip_file(final_zip_save_path, EXTRACT_DIR):
                print(f"文件 {zip_filename} 校验通过但解压失败.")
        else:
            print(f"文件 {zip_filename} 校验和验证失败. 不进行解压.")
            # 可选: 删除校验失败的zip文件以供下次重新下载
            # os.remove(final_zip_save_path)

        # 5. 清理CHECKSUM文件
        if os.path.exists(checksum_save_path):
            try:
                os.remove(checksum_save_path)
                print(f"已删除CHECKSUM文件: {checksum_save_path}")
            except OSError as e:
                print(f"无法删除CHECKSUM文件 {checksum_save_path}: {e}")

        time.sleep(1)  # 对服务器友好, 每次循环后等待1秒

        # 更新月份和年份
        current_month += 1
        if current_month > 12:
            current_month = 1
            current_year += 1
            if current_year > end_year:
                break

    print("\n所有任务已完成.")


if __name__ == "__main__":
    main()